# Copyright (C) 2009  Rob Britton
#
# This file is part of OORegress.
#
# OORegress is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

###############################
#
# Some handy stats functions
#
###############################

require 'parser'
require 'matrix'
require 'statistics2'

VARIABLES_OFFSET = 14
ANOVA_OFFSET = 8
BASIC_OFFSET = 1

class DataColumn
  attr_accessor :name
  attr_reader :data

  def initialize name, data = nil
    @data = data || []
    @name = name
  end

  def size; @data.size; end
  def <<(x); @data << x; end
  def [](i); @data[i]; end

  def to_s
    "[#{@name}: #{@data.join(", ")}]"
  end
end

class DataSet
  attr_reader :data

  def initialize data = nil
    @data = data || {}
  end

  def num_vars
    @data.size
  end

  def num_obs
    @data.map { |label, values| values.size }.min
  end

  def [](label)
    @data[label]
  end

  def <<(column)
    raise "Must pass a DataColumn to DataSet.<<" unless column.is_a?(DataColumn)

    @data[column.name] = column
  end

  def to_s
    @data.values.map { |d| d.to_s }.join("\n")
  end
end

def load_dataset calc
  dataset = DataSet.new

  i = 0
  while "" != (label = calc[i, 0])
    # found a column
    dc = DataColumn.new label

    j = 1
    while "" != (value = calc[i, j])
      dc << value
      j += 1
    end

    dataset << dc
    i += 1
  end

  dataset
end

def regress calc, *params
  lhs, rhs = parse_eq params.first

  # fetch the labels
  data = load_dataset calc
  Identifier.dataset = data

  intercept = !params.include?("no intercept")

  num_vars = rhs.size
  num_observations = data.num_obs

  lhs_min_obs = lhs.min_obs
  rhs_min_obs = rhs.min_obs

  lhs_values = []
  rhs_values = Array.new(num_vars) { [] }
  
  min_obs = lhs_min_obs > rhs_min_obs ? lhs_min_obs : rhs_min_obs
  (min_obs..(num_observations - 1)).each do |i|
    lhs_values << lhs.value(i)
    rhs.value_array(i).each_with_index do |v, j|
      rhs_values[j] << v
    end
  end

  num_observations -= min_obs
  
  if intercept
    rhs_values.unshift([1.0] * num_observations)
    num_vars += 1
  end

  x = Matrix.columns(rhs_values)
  y = Matrix.column_vector(lhs_values)
  xt = x.t
  
  betas = (xt * x).inverse * xt * y     # regression coefficients
  yhat = x * betas                      # lhs estimates
  ehat = y - yhat                       # residuals

  err_ss = (ehat.t * ehat)[0, 0]        # sum of squared residuals
  dof = num_observations - num_vars     # degrees of freedom
  ymean = y.to_a.inject(0.0) { |s,a| s + a.first.to_f } / num_observations.to_f
  tot_ss = y.to_a.inject(0.0) { |s,a| s + (a.first.to_f - ymean)**2 }  # total sum of squares
  r_sq = 1.0 - err_ss / tot_ss          # R^2 value

  means = xt.to_a.map { |r|             # means of rhs variables
    r.inject(0) { |s,a| s + a } / num_observations
  }
  sqs = xt.to_a.zip(means).map { |r, m| # squared distance from the mean of the rhs variables
    r.inject(0) { |s,a| s + (a - m)**2 }
  }

  # create a new sheet for the regression output
  target_sheet =
    if params.find { |p| p =~ /^output (.*)$/ }
      $1
    else
      "Regression Results"
    end

  calc.insertSheet target_sheet

  # basic info
  calc[0, BASIC_OFFSET] = "Basic Information"
  calc.setStyle(0, BASIC_OFFSET, "Result")
  %w[Multiple-R R-Squared Adjusted\ R-Squared Standard\ Error Observations].each_with_index do |l, i|
    calc[0, BASIC_OFFSET + i + 1] = l
  end

  [
    Math.sqrt(r_sq),
    r_sq,
    1 - (1 - r_sq) * (num_observations - 1) / dof,
    Math.sqrt(err_ss / dof),
    num_observations
  ].each_with_index do |v, i|
    calc[1, BASIC_OFFSET + i + 1] = v
  end

  # ANOVA
  calc[0, ANOVA_OFFSET] = "ANOVA"
  calc.setStyle(0, ANOVA_OFFSET, "Result")
  %w[Regression Residual Total].each_with_index do |l, i|
    calc[0, ANOVA_OFFSET + 2 + i] = l
  end
  %w[dof SS MS F Significance\ F].each_with_index do |l, i|
    calc[1 + i, ANOVA_OFFSET + 1] = l
  end

  # degrees of freedom
  calc[1, ANOVA_OFFSET + 2] = num_vars - 1
  calc[1, ANOVA_OFFSET + 3] = dof
  calc[1, ANOVA_OFFSET + 4] = num_observations - 1

  # sum of squares
  calc[2, ANOVA_OFFSET + 2] = tot_ss - err_ss
  calc[2, ANOVA_OFFSET + 3] = err_ss
  calc[2, ANOVA_OFFSET + 4] = tot_ss

  # MS
  calc[3, ANOVA_OFFSET + 2] = tot_ss - err_ss / (num_vars - 1)
  calc[3, ANOVA_OFFSET + 3] = err_ss / dof

  # F
  f = (r_sq / (1.0 - r_sq)) * dof.to_f / (num_observations - 1).to_f
  calc[4, ANOVA_OFFSET + 2] = f
  calc[5, ANOVA_OFFSET + 2] = 1.0 - Statistics2.fdist(dof, num_vars - 1, f)

  # info on variables
  calc[0, VARIABLES_OFFSET] = "Coefficients"
  calc.setStyle(0, VARIABLES_OFFSET, "Result")

  %w[Coefficients Standard\ Error T-stat p-value Lower\ 95% Upper\ 95%].each_with_index do |l, i|
    calc[i + 1, VARIABLES_OFFSET + 1] = l
  end

  calc[0, VARIABLES_OFFSET + 2] = "Intercept" if intercept

  rhs.text_array.each_with_index do |l, i|
    calc[0, VARIABLES_OFFSET + i + (intercept ? 3 : 2)] = l
  end

  t95 = Statistics2.ptdist(dof, 0.025)

  num_vars.times do |i|
    coeff = betas[i, 0]
    stderr = Math.sqrt(err_ss / (dof * sqs[i]))
    t = coeff / stderr
    calc[1, VARIABLES_OFFSET + i + 2] = coeff                                 # coefficient
    calc[2, VARIABLES_OFFSET + i + 2] = stderr                                # standard error
    calc[3, VARIABLES_OFFSET + i + 2] = t                                     # t-stat
    calc[4, VARIABLES_OFFSET + i + 2] = Statistics2.tdist(dof, -t.abs)        # p-value
    calc[5, VARIABLES_OFFSET + i + 2] = coeff - t95*stderr                    # lower 95%
    calc[6, VARIABLES_OFFSET + i + 2] = coeff + t95*stderr                    # upper 95%
  end

  if params.include?("residuals")
    res_offset = VARIABLES_OFFSET + 3 + num_vars

    calc[0, res_offset] = "Residuals"
    calc.setStyle(0, res_offset, "Result")
    ["Observation", "Estimated #{lhs.text[1..-2]}", "Residual"].each_with_index do |l, i|
      calc[i, res_offset + 1] = l
    end

    ehat.to_a.each_with_index do |e, i|
      calc[0, res_offset + 2 + i] = i + 1
      calc[1, res_offset + 2 + i] = yhat[i, 0]
      calc[2, res_offset + 2 + i] = e.first
    end
  end
end

def generate_randoms(calc, *params)
  name = ""
  if params.find { |p| p =~ /^named ([a-z][a-z0-9_]*)$/i }
    name = $1
  else
    puts "Please enter a name for your random variable." and return
  end

  num = 0
  if params.find { |p| p =~ /^n (\d+)$/ }
    num = $1.to_i
  else
    puts "Please input the number of values you would like." and return
  end

  pos = [0, 0]

  if (col = params.find { |p| p =~ /^at ([a-z]*)(\d*)$/i })
    pos[0] = $1[0] - 65 if $1 != ""
    pos[1] = $2.to_i - 1 if $2 != ""
  else
    # find the first empty column
    i = 0
    while "" != calc[i, 0]
      i += 1
    end
    pos[0] = i
  end

  calc[pos[0], pos[1]] = name

  # generate the numbers
  distribution =
    if params.find { |p| p =~ /^distributed ([a-z]+)$/i }
      $1
    else
      "uniformly"
    end

  case distribution
  when "uniformly"
    integer = params.map { |p| p.downcase }.include?("integer")

    bounds =
      if integer
        [$1.to_i, $2.to_i].sort if params.find { |p| p =~ /^between (\-?\d+) and (\-?\d+)$/ }
      else
        [$1.to_f, $3.to_f].sort if params.find { |p| p =~ /^between (\-?\d+(\.\d+)?) and (\-?\d+(\.\d+)?)$/ }
      end
    bounds ||= [0, 1]

    num.times do |i|
      calc[pos[0], pos[1] + i + 1] =
        if integer
          rand(bounds[1] - bounds[0]) + bounds[0]
        else
          rand * (bounds[1] - bounds[0]) + bounds[0]
        end
    end
  when "normally"
    mean =
      if params.find { |p| p =~ /^mean (\d+(\.\d+)?)$/ }
        $1.to_f
      else
        0.0
      end
    stdev =
      if params.find { |p| p =~ /^stdev (\d+(\.\d+)?)$/ }
        $1.to_f
      else
        1.0
      end

    num.times do |i|
      rnd1, rnd2 = rand, rand
      calc[pos[0], pos[1] + i + 1] = mean + stdev * Math.sqrt(-2 * Math.log(rnd1)) * Math.cos(6.28 * rnd2)
    end
  end
end
